package cn.genmer.test.security.machinelearning.deeplearning4j.begin;

import org.deeplearning4j.datasets.iterator.impl.MnistDataSetIterator;
import org.deeplearning4j.nn.api.OptimizationAlgorithm;
import org.deeplearning4j.nn.conf.MultiLayerConfiguration;
import org.deeplearning4j.nn.conf.NeuralNetConfiguration;
import org.deeplearning4j.nn.conf.layers.DenseLayer;
import org.deeplearning4j.nn.conf.layers.OutputLayer;
import org.deeplearning4j.nn.multilayer.MultiLayerNetwork;
import org.deeplearning4j.optimize.listeners.ScoreIterationListener;
import org.deeplearning4j.util.ModelSerializer;
import org.nd4j.linalg.activations.Activation;
import org.nd4j.linalg.dataset.api.iterator.DataSetIterator;
import org.nd4j.linalg.factory.Nd4j;
import org.nd4j.linalg.lossfunctions.LossFunctions;

import java.io.File;
import java.io.IOException;

public class Deeplearning4jTraining {

    public static void main(String[] args) throws IOException {
        int numInputs = 784;
        int numOutputs = 10;
        //
        int batchSize = 64;
        // 训练次数
        int numEpochs = 1;

        // 设置 Nd4j 的后端为 CPU
        Nd4j.getMemoryManager().setAutoGcWindow(10000);
//        Nd4j.getBackend().setEnableDebugMode(true);
//        Nd4j.getBackend().setEnableVerboseMode(true);

        // 加载 MNIST 数据集
        DataSetIterator mnistTrain = new MnistDataSetIterator(batchSize, true, 12345);

        // 构建神经网络配置
        MultiLayerConfiguration conf = new NeuralNetConfiguration.Builder()
                .seed(12345)
                .optimizationAlgo(OptimizationAlgorithm.STOCHASTIC_GRADIENT_DESCENT)
                .list()
                .layer(0, new DenseLayer.Builder().nIn(numInputs).nOut(256).build())
                .layer(1, new OutputLayer.Builder(LossFunctions.LossFunction.NEGATIVELOGLIKELIHOOD)
                        .nIn(256).nOut(numOutputs).activation(Activation.SOFTMAX).build())
                .build();

        // 构建多层感知器模型
        MultiLayerNetwork model = new MultiLayerNetwork(conf);
        model.init();
        model.setListeners(new ScoreIterationListener(10));

        // 训练模型
        for (int i = 0; i < numEpochs; i++) {
            model.fit(mnistTrain);
        }

        // 保存训练好的模型
        File modelFile = new File("/Users/genmer/Documents/Codes/tensorFlowModel/deepLearning4jModel.zip");
        ModelSerializer.writeModel(model, modelFile, true);
    }
}